#include "airkiss.h"
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <errno.h>
#include <unistd.h>
#include <limits.h>
#include <syslog.h>

#include "ieee80211.h"

#if DEBUG
#define airkiss_log(x, ...) printf(__VA_ARGS__)
#else
#define airkiss_log    syslog
#endif

uint8_t CRC8(uint8_t* message, uint8_t len);

/*
 * AirKiss 流程
0. 获取当前无线环境中所有非隐蔽AP的ssid，rssi以及信道(非必要的过程,但是可以节省配网时间)
1. 切换信道
2. 等待前导数据，确定长度基数,前导数据为4个连续的数组成，用于确定基数，重复多次
3. 如果接收前导数据超时，则跳转到1，否则跳转到4
4. 接收magic code字段，重复多次
magic code 字段定义
               | 高5位  | 低4位           |
第  1 个 9bits | 0x0    | length(high)    | (当总长度小于16时，这里最高位会被设置了1,这里应该是避免广播包长度为0的可能)
第  2 个 9bits | 0x1    | length(low)     |
第  3 个 9bits | 0x2    | ssid crc8(high) |
第  4 个 9bits | 0x3    | ssid crc8(low)  |
解析得到数据长度，和ssid的 CRC值（可以加速陪网过程，先不实现）
5. 接收prefix code字段,重复多次
               | 高5位  | 低4位              |
第  1 个 9bits | 0x4    | pwd length(high)   |
第  2 个 9bits | 0x5    | pwd length(low)    |
第  3 个 9bits | 0x6    | pwd len crc8(high) |
第  4 个 9bits | 0x7    | pwd len crc8(low)  |
6. 接收数据Sequence，数据sequence包括两部分: sequence header和data, 重复多次
Sequence header字段
              | 8 | 7 | 6 ~ 0
第  1 个 9bits | 0 | 1 | sequnece crc8(low 7bits) |
第  2 个 9bits | 0 | 1 | sequnece index           |
data 字段
              | 8 | 7 ~ 0
第  1 个 9bits | 1 | data |
第  2 个 9bits | 1 | data |
第  3 个 9bits | 1 | data |
第  4 个 9bits | 1 | data |
*/
#define AIRKISS_LEADING_SIZE  4
#define AIRKISS_SEQUENCE_SIZE 6
typedef struct {
    unsigned short seqno;
    unsigned short data;
} __airkiss_data;
enum {
    AIRKISS_INIT = 0,
    AIRKISS_LEADING,
    AIRKISS_LEADING_FIN,
    AIRKISS_MAGICCODE,
    AIRKISS_MAGIC_FIN,
    AIRKISS_PREFIXCODE,
    AIRKISS_PREFIX_FIN,
    AIRKISS_SEQUENCE,
    AIRKISS_DONE
};
struct __airkiss_connection {

    int state;

    // addr2 + addr3
    unsigned char key[12];
    // 组包临时变量
    __airkiss_data tmp[6];
    unsigned short tpos;

    // 前导数据
    uint16_t basecode;
    uint8_t preflags;
    // Magic code
    uint8_t datalen;
    uint8_t ssidcrc;
    uint8_t seqsz;
    // Prefix Code
    uint8_t pwdlength;
    uint8_t pwdlencrc;
    // Result
    uint8_t result[32+1+64];
    uint8_t resmask[25]; // (32+1+64) / 4 = 25

    // 组数据的临时缓冲
    //__airkiss_data tdata[150]; // = ((32+1+64) / 4) * 6 = 150
} __attribute__((__packed__));
typedef struct __airkiss_connection __airkiss_connection;


void _waiting_leading(__airkiss_connection* c, unsigned short seq, unsigned short d) {
    unsigned short l;
    if (c->tmp[c->tpos - 1].seqno == seq) {
        // 重复数据
        return ;
    }
    c->tmp[c->tpos].data = d;
    c->tmp[c->tpos].seqno = seq;
    c->tpos++;
    if (c->tpos >= AIRKISS_LEADING_SIZE) {
        // 需要4个连续递增的数据
        seq = c->tmp[0].seqno;
        d = c->tmp[0].data;
        for (l = 1; l < (AIRKISS_LEADING_SIZE); l++) {
            seq++;d++;
            //airkiss_log("i%u: seq=%u|%u, data=%u|%u\n", l, seq, c->tmp[l].seqno, d, c->tmp[l].data);
            if (seq != c->tmp[l].seqno || d != c->tmp[l].data)
                break;
        }

        if (l == AIRKISS_LEADING_SIZE) {
            c->tpos = 0;
            c->basecode = c->tmp[0].data - 1;
            c->state = AIRKISS_LEADING_FIN;
            airkiss_log(LOG_DEBUG, "Change state to %d\nbasecode=%u\n", c->state, c->basecode);
        }
        else {
            c->tmp[0] = c->tmp[1];
            c->tmp[1] = c->tmp[2];
            c->tmp[2] = c->tmp[3];
            c->tpos = 3;
        }
    }
}

void _waiting_magiccode(__airkiss_connection* c, unsigned short seq, unsigned short d) {
/*
magic code 字段定义
               | 高5位  | 低4位           |
第  1 个 9bits | 0x0    | length(high)    |
第  2 个 9bits | 0x1    | length(low)     |
第  3 个 9bits | 0x2    | ssid crc8(high) |
第  4 个 9bits | 0x3    | ssid crc8(low)  |
解析得到数据长度，和ssid的 CRC值（可以加速陪网过程，先不实现）
*/
    (void) seq;
    unsigned short idx = d & 0x1f0;
    switch (idx) {
    case 0x00:
        d &= 0x0f;
        c->datalen |= (d << 4);
        c->preflags |= (1 << 0);
        break;
    case 0x10:
        d &= 0x0f;
        c->datalen |= (d);
        c->preflags |= (1 << 1);
        break;
    case 0x20:
        d &= 0x0f;
        c->ssidcrc |= (d << 4);
        c->preflags |= (1 << 2);
        break;
    case 0x30:
        d &= 0x0f;
        c->ssidcrc |= (d);
        c->preflags |= (1 << 3);
        break;
    default:
        break;
    }
    if (0x0f == c->preflags) {
        c->datalen &= 0x7f;
        c->seqsz = c->datalen + 3;
        c->seqsz >>= 2;
        c->state = AIRKISS_MAGIC_FIN;
        airkiss_log(LOG_DEBUG, "Change state to %d\ndatalen=%u, ssidcrc=%u, seqsz=%u\n", c->state, c->datalen, c->ssidcrc, c->seqsz);
    }
}

void _waiting_prefix(__airkiss_connection* c, unsigned short seq, unsigned short d) {
/*
              | 高5位   | 低4位           |
第  1 个 9bits | 0x4    | pwd length(high)   |
第  2 个 9bits | 0x5    | pwd length(low)    |
第  3 个 9bits | 0x6    | pwd len crc8(high) |
第  4 个 9bits | 0x7    | pwd len crc8(low)  |
*/
    (void) seq;
    unsigned short idx = d & 0x1f0;
    switch (idx) {
    case 0x40:
        d &= 0x0f;
        c->pwdlength |= (d << 4);
        c->preflags |= (1 << 4);
        break;
    case 0x50:
        d &= 0x0f;
        c->pwdlength |= (d);
        c->preflags |= (1 << 5);
        break;
    case 0x60:
        d &= 0x0f;
        c->pwdlencrc |= (d << 4);
        c->preflags |= (1 << 6);
        break;
    case 0x70:
        d &= 0x0f;
        c->pwdlencrc |= (d);
        c->preflags |= (1 << 7);
        break;
    default:
        break;
    }
    if (0xf0 == c->preflags) {
        // CRC check
        uint8_t crc = CRC8(&c->pwdlength, 1);
        if (crc == c->pwdlencrc) {
            c->state = AIRKISS_PREFIX_FIN;
            airkiss_log(LOG_DEBUG, "Change state to %d\npwdlen=%u, pwdlencrc=%u\n", c->state, c->pwdlength, c->pwdlencrc);
        }
    }
}

int _datacalc(__airkiss_data* tmp, uint8_t len, uint8_t dlen, uint8_t* res) {
    uint8_t crcbuf[5];
    if (len < dlen + 2) // 数据长度不足
        return -1;
    int ok = 1;
    unsigned short s = tmp[0].seqno;
    for (uint8_t i = 1; i < dlen + 2; i++) {
        s++; if (s >= 4096) s = 0;
        if (s != tmp[i].seqno) {
            ok = 0;
            break;
        }
    }
    if (!ok) { // 数据不连续
        airkiss_log(LOG_ERR, "\tNot continous data\n");
        return -2;
    }
    // CRC checking
    for (uint8_t i = 1; i < dlen + 2; i++) {
        crcbuf[i - 1] = tmp[i].data & 0xff;
    }
    crcbuf[0] &= 0x7f;
    uint8_t crc = CRC8(crcbuf, dlen + 1);
    if ((crc & 0x7f) != (tmp[0].data & 0x7f)) { // CRC校验失败
        airkiss_log(LOG_ERR, "\tCRC8 error\n");
        return -3;
    }
    // 放到结果中
    memcpy(res, &crcbuf[1], dlen);
    return 0;
}
static inline void _datapush2tmp(__airkiss_connection* c) {
    /*airkiss_log(LOG_DEBUG, "[index][%u] [%04u %04u %04u %04u %04u %04u]\n"
                "[ data][%u] [ %03x  %03x  %03x  %03x  %03x  %03x]\n",
                c->tpos,
                c->tmp[0].seqno, c->tmp[1].seqno, c->tmp[2].seqno, c->tmp[3].seqno, c->tmp[4].seqno, c->tmp[5].seqno,
                c->tpos,
                c->tmp[0].data & 0x1ff, c->tmp[1].data & 0x1ff, c->tmp[2].data & 0x1ff,
                c->tmp[3].data & 0x1ff, c->tmp[4].data & 0x1ff, c->tmp[5].data & 0x1ff);*/
    if ((c->tmp[0].seqno + 1) != c->tmp[1].seqno) {// Not continous sequence number
        //airkiss_log(LOG_DEBUG, "Error sequence header\n");
        return ;
    }
    unsigned short idx = (c->tmp[1].data & 0x7f);
    if (idx >= c->seqsz) { // Wrong data sequence number
        //airkiss_log(LOG_DEBUG, "Error data sequence no %u|%u\n", idx, c->seqsz);
        return ;
    }
    uint8_t dlen = 4;
    if ((idx == (c->seqsz - 1)) && (c->datalen & 0x03)) {
        dlen = c->datalen & 0x03;
    }
    if (1 == c->resmask[idx]) {
        //airkiss_log(LOG_DEBUG, "sequence index %u is set\n", idx);
    }
    else if (0 == _datacalc(c->tmp, c->tpos, dlen, &(c->result[idx*4]))) {
        c->resmask[idx] = 1;
        airkiss_log(LOG_DEBUG, "[result][%u][%u] [%02x %02x %02x %02x]\n",
               idx, dlen, c->result[idx*4], c->result[idx*4+1], c->result[idx*4+2], c->result[idx*4+3]);
    }
    else {
#if 0
        airkiss_log(LOG_DEBUG, "Incompeted data\n");
        /* 缓存数据
         * 实际上收到的airkiss数据帧SEQUENCE并不能保证是连续的，
         * 这似乎是取决于发包端的，这给重组包带来了很大的麻烦
         * N: 本来算法期望是每次抓到不同位置的包，最终通过帧号判断串联起来成为完整数据，
         *    但调试结果显示当丢包时每次丢的包都是一个位置，这段实际上没有起到什么作用
         */
        __airkiss_data* d = &(c->tdata[idx * 6]);
        __airkiss_data* t = c->tmp;
        d[0] = t[0]; d[1] = t[1];
        for (unsigned short i = 2; i < c->tpos; i++) {
            unsigned short ii = t[i].seqno - t[0].seqno;
            if (ii < dlen) {
                d[ii] = t[i];
            }
            else {
                break;
            }
        }
        // 尝试解析
        uint8_t crcbuf[5];
        // CRC checking
        for (uint8_t i = 1; i < dlen + 2; i++) {
            crcbuf[i - 1] = d[i].data & 0xff;
        }
        crcbuf[0] &= 0x7f;
        uint8_t crc = CRC8(crcbuf, dlen + 1);
        if ((crc & 0x7f) == (d[0].data & 0x7f)) { // CRC校验成功
            // 放到结果中
            memcpy(&(c->result[idx*4]), &crcbuf[1], dlen);
            airkiss_log(LOG_DEBUG, "[FIXED result] [%u] [%02x %02x %02x %02x]\n",
                        dlen, crcbuf[1], crcbuf[2], crcbuf[3], crcbuf[4]);
            c->resmask[idx] = 1;
        }
#endif
    }
    // 判断结束状态
    int ok = 1;
    for (uint8_t i = 0; i < c->seqsz; i++) {
        if (0 == c->resmask[i]) {
            ok = 0;
            break;
        }
    }
    if (ok) {
        c->state = AIRKISS_DONE;
    }
}
void _waiting_sequence(__airkiss_connection* c, unsigned short seq, unsigned short d) {
    /*
    Sequence header字段
                  | 8 | 7 | 6 ~ 0
    第  1 个 9bits | 0 | 1 | sequnece crc8(low 7bits) |
    第  2 个 9bits | 0 | 1 | sequnece index           |
    data 字段
                  | 8 | 7 ~ 0
    第  1 个 9bits | 1 | data |
    第  2 个 9bits | 1 | data |
    第  3 个 9bits | 1 | data |
    第  4 个 9bits | 1 | data |
    */
    //airkiss_log(LOG_DEBUG, "%04u|%03x\n", seq, d & 0x1ff);
    if (seq == c->tmp[c->tpos - 1].seqno) {
        // 不知道为何抓到重复数据
        c->tmp[c->tpos - 1].data = d;
    }
    else if (0x80 == (d & 0x180)) {
        // Sequence header
        if (c->tpos > 2) {
            _datapush2tmp(c);
            c->tpos = 0;
            memset(c->tmp, 0, sizeof(c->tmp));
        }
        c->tmp[c->tpos].data = d;
        c->tmp[c->tpos].seqno = seq;
        c->tpos ++;
    }
    else if (d & 0x100) {
        // Sequence data
        if (c->tpos < 2) {
            c->tpos = 0;
        }
        else if (seq - c->tmp[0].seqno >= 6) {
            _datapush2tmp(c);
            c->tpos = 0;
            memset(c->tmp, 0, sizeof(c->tmp));
        }
        else {
            c->tmp[c->tpos].data = d;
            c->tmp[c->tpos].seqno = seq;
            c->tpos ++;
        }
    }
    else {
        _datapush2tmp(c);
        c->tpos = 0;
        memset(c->tmp, 0, sizeof(c->tmp));
    }
}

int airkiss_input(__airkiss_connection* c, unsigned short seq, unsigned short payloadlen)
{
    unsigned short data = 0;
    switch (c->state) {
    case AIRKISS_INIT:
        c->tpos = 0;
        c->state = AIRKISS_LEADING;
    case AIRKISS_LEADING:
        _waiting_leading(c, seq, payloadlen);
        break;
    case AIRKISS_LEADING_FIN:
        data = payloadlen - c->basecode;
        if (data <= 4)
            break;
        else {
            c->preflags = 0;
            c->state = AIRKISS_MAGICCODE;
            airkiss_log(LOG_DEBUG, "Change state to %d\n", c->state);
        }
    case AIRKISS_MAGICCODE:
        data = payloadlen - c->basecode;
        _waiting_magiccode(c, seq, data);
        break;
    case AIRKISS_MAGIC_FIN:
    {
        data = payloadlen - c->basecode;
        unsigned short idx = data & 0x1f0;
        idx >>= 4;
        if (idx <= 0x3)
            break;
        else {
            c->tpos = 0;
            c->preflags = 0;
            c->state = AIRKISS_PREFIXCODE;
            airkiss_log(LOG_DEBUG, "Change state to %d\n", c->state);
        }
    }
    case AIRKISS_PREFIXCODE:
        data = payloadlen - c->basecode;
        _waiting_prefix(c, seq, data);
        break;
    case AIRKISS_PREFIX_FIN:
    {
        data = payloadlen - c->basecode;
        unsigned short idx = data & 0x1f0;
        idx >>= 4;
        if (idx <= 0x7)
            break;
        else {
            c->tpos = 0;
            c->state = AIRKISS_SEQUENCE;
            airkiss_log(LOG_DEBUG, "Change state to %d\n", c->state);
        }
    }
    case AIRKISS_SEQUENCE:
        data = payloadlen - c->basecode;
        _waiting_sequence(c, seq, data);
        break;
    case AIRKISS_DONE:
        _waiting_sequence(c, seq, data);
        break;
    default:
        break;
    }

    return 0;
}

uint8_t CRC8(uint8_t* message, uint8_t len)
{
    uint8_t crc = 0;
    uint8_t i;

    while (len--) {
        crc ^= *message++;
        for (i = 0; i < 8; i++) {
            if (crc & 0x01)
                crc = (crc >> 1) ^ 0x8c;
            else
                crc >>= 1;
        }
    }

    return crc;
}

typedef struct {
#if (__WORDSIZE == 64)
    airkiss_memset_fn memset;
	airkiss_memcpy_fn memcpy;
	airkiss_memcmp_fn memcmp;
	airkiss_printf_fn printf;
    int rev[18];

    __airkiss_connection* connection[16];
#elif (__WORDSIZE == 32)
    airkiss_memset_fn memset;
	airkiss_memcpy_fn memcpy;
	airkiss_memcmp_fn memcmp;
	airkiss_printf_fn airkiss_log;
    int rev[22];

    __airkiss_connection* connection[32];
#else
    int dummyap[26];
    int dummy[32];
#endif
} __airkiss_context_t;
#define _CNNSIZE (sizeof(ctx->connection)/sizeof(ctx->connection[0]))

int airkiss_init(airkiss_context_t* context, const airkiss_config_t* config)
{
    airkiss_log(LOG_DEBUG, "airkiss_context_t init");
    __airkiss_context_t* ctx = (__airkiss_context_t*) context;
    if (config) {
        ctx->memset = config->memset;
        ctx->memcmp = config->memcmp;
        ctx->memcpy = config->memcpy;
        ctx->printf = config->printf;
    }
    for (int i = 0; i < _CNNSIZE; i++) {
        ctx->connection[i] = (__airkiss_connection*) malloc(sizeof(__airkiss_connection));
        ctx->memset(ctx->connection[i], 0, sizeof(__airkiss_connection));
    }
    return 0;
}

static char _airkiss_version[] = "1";
const char* airkiss_version(void)
{
    return _airkiss_version;
}

int airkiss_recv(airkiss_context_t* context, const void* frame, unsigned short length)
{
    uint8_t zkey[12] = {0};
    unsigned char key[12];
    int status = AIRKISS_INIT;
    __airkiss_context_t* ctx = (__airkiss_context_t*) context;
    if (length < sizeof(mac_t)) {
        goto status_checking;
    }
    mac_t* mac = (mac_t*) frame;

    if (IEEE80211_FTYPE_DATA == mac->type) {
        if (mac->subtype & IEEE80211_STYPE_NULLFUNC) {
            goto status_checking;
        }
        // Accept fromDS=1 && toDS=0 or fromDS=0 && toDS=1
        // Discard fromDS=0 && toDS=0 or fromDS=1 && toDS=1
        if (mac->from_ds == mac->to_ds) {
            goto status_checking;
        }
        unsigned short payloadlen = length - sizeof(mac_t);
        if (mac->subtype & IEEE80211_STYPE_QOS_DATA) {
            payloadlen -= QOS_SIZE;
        }
        if (payloadlen < LLC_SIZE)
            goto status_checking;
        payloadlen -= LLC_SIZE;
        // connection
        memcpy(key, mac->addr2, 6);
        memcpy(key + 6, mac->addr3, 6);
        for (unsigned int i = 0; i < _CNNSIZE; i++) {
            if ((memcmp(ctx->connection[i]->key, zkey, sizeof(zkey)) == 0 && memcpy(ctx->connection[i]->key, key, sizeof(key)))
                     || memcmp(ctx->connection[i]->key, key, sizeof(key)) == 0) {
                //;
                unsigned short seq = mac->sequence;
                seq >>= 4;
                airkiss_input(ctx->connection[i], seq, payloadlen);
                break;
            }
        }
    }
    // Status checking
status_checking:
    for (unsigned int i = 0; i < _CNNSIZE; i++) {
        if (ctx->connection[i]->state > status) {
            status = ctx->connection[i]->state;
        }
    }
    if (AIRKISS_DONE == status)
        return AIRKISS_STATUS_COMPLETE;
    else if (status > AIRKISS_LEADING)
        return AIRKISS_STATUS_CHANNEL_LOCKED;
    return AIRKISS_STATUS_CONTINUE;
}

int airkiss_get_result(airkiss_context_t* context, airkiss_result_t* result)
{
    if (!context || !result) {
        return -1;
    }
    __airkiss_context_t* ctx = (__airkiss_context_t*) context;
    for (unsigned int i = 0; i < _CNNSIZE; i++) {
        if (ctx->connection[i]->state == AIRKISS_DONE) {
            //
            result->pwd = (char*)ctx->connection[i]->result;
            result->pwd_length = ctx->connection[i]->pwdlength;
            result->random = ctx->connection[i]->result[result->pwd_length];
            result->ssid = (char*)&ctx->connection[i]->result[result->pwd_length + 1];
            result->ssid_length = ctx->connection[i]->datalen - ctx->connection[i]->pwdlength - 1;
            return 0;
        }
    }
    return -1;
}

int airkiss_change_channel(airkiss_context_t* context)
{
    __airkiss_context_t* ctx = (__airkiss_context_t*) context;
    for (unsigned int i = 0; i < _CNNSIZE; i++) {
        ctx->memset(ctx->connection[i], 0, sizeof(__airkiss_connection));
    }
    return 0;
}

#if AIRKISS_ENABLE_CRYPT
int airkiss_set_key(airkiss_context_t* context, const unsigned char* key, unsigned int length)
{
    return -10;
}
#endif
